#include <cassert>
#include <chrono>
#include <climits>
#include <iostream>
#include <map>
#include <random>
#include <vector>
using namespace std;

const int n = 1e7;
template <typename KeyType, typename ValueType> struct SkipListNode {
  int level;       // in redis, is level
  KeyType key;     // in redis, is score
  ValueType value; // in redis, is ele
  SkipListNode **forward;

  SkipListNode() {}

  SkipListNode(KeyType k, ValueType v, int l, SkipListNode *next = NULL) {
    key = k;
    value = v;
    level = l;
    forward = new SkipListNode *[l + 1];
    for (int i = 0; i <= l; ++i)
      forward[i] = next;
  }

  // in redis: zslFreeNode
  ~SkipListNode() {
    if (forward != NULL)
      delete[] forward;
  }
};

template <typename KeyType, typename ValueType> struct SkipList {
  // Maximum level for this skip list
  int maxLevel;
  // Probability factor for random level generation
  int probabilityFactor;
  // Maximum value for random number generation
  int maxValue;

  // Threshold for determining level based on probability
  int probabilityThreshold;
  // Sentinel value for invalid keys
  static constexpr int INVALID_KEY = INT_MAX;

  // to store rand generator
  vector<int> *randVec = NULL;
  long randPointer = 0;
  static const int RAND_NUMBER = n;
  //

  // Constructor with default parameters
  SkipList(int maxLevel = 32, int probabilityFactor = 4, int maxValue = 0xFFFF)
      : maxLevel(maxLevel), probabilityFactor(probabilityFactor),
        maxValue(maxValue) {
    probabilityThreshold = maxValue / probabilityFactor;

    std::random_device rd;
    randomGenerator = std::mt19937(rd());
    distribution = std::uniform_int_distribution<>(0, maxValue);

    currentLevel = length = 0;
    tail = new SkipListNode<KeyType, ValueType>(INVALID_KEY, 0, 0);
    head = new SkipListNode<KeyType, ValueType>(INVALID_KEY, 0, maxLevel, tail);

    // add rand gen
    randVec = new vector<int>(RAND_NUMBER);
    for (int i = 0; i < RAND_NUMBER; ++i) {
      (*randVec)[i] = distribution(randomGenerator);
    }
    //
  }

  SkipListNode<KeyType, ValueType> *head, *tail;
  // total skipList length
  int length;
  int currentLevel;

  ~SkipList() {
    delete head;
    delete tail;
    if (randVec != NULL)
      delete randVec;
  }

  // in redis: zslRandomLevel
  int generateRandomLevel_orig() {
    int level = 1;
    // in redis equal to: (random()&0xFFFF) < (ZSKIPLIST_P * 0xFFFF)
    while (distribution(randomGenerator) < probabilityThreshold)
      ++level;
    // in redis maxLevel is restricted to ZSKIPLIST_MAXLEVEL, which is 32
    return maxLevel > level ? level : maxLevel;
  }

  // optimized version
  int generateRandomLevel() {
    int level = 1;
    // in redis equal to: (random()&0xFFFF) < (ZSKIPLIST_P * 0xFFFF)
    while ((*randVec)[randPointer++ % RAND_NUMBER] < probabilityThreshold)
      ++level;
    // in redis maxLevel is restricted to ZSKIPLIST_MAXLEVEL, which is 32
    return maxLevel > level ? level : maxLevel;
  }

  // in redis: zslCreateNode
  SkipListNode<KeyType, ValueType> *
  createNode(const KeyType &key, const ValueType &value, int level) {
    return new SkipListNode<KeyType, ValueType>(key, value, level);
  }

  // in redis: zslInsert / zslUpdateScore
  // zslInsert assume element not already inside the skiplist, and redis allows
  // skipList has same key(score) but different value(ele), so it commands its
  // caller check element is in skiplist first. If so, update. zslUpdateScore,
  // zslInsert official caller, use to update score. insert combines two.
  void insert(const KeyType &key, const ValueType &value) {
    SkipListNode<KeyType, ValueType> *update[maxLevel + 1];

    SkipListNode<KeyType, ValueType> *current = head;
    for (int i = currentLevel; i >= 0; --i) {
      while (current->forward[i]->key < key) {
        current = current->forward[i];
      }
      update[i] = current;
    }
    current = current->forward[0];

    if (current->key == key) {
      current->value = value;
      return;
    }

    int newLevel = generateRandomLevel();
    if (newLevel > currentLevel) {
      newLevel = ++currentLevel;
      update[newLevel] = head;
    }

    auto *newNode = createNode(key, value, newLevel);
    for (int i = newLevel; i >= 0; --i) {
      // in redis:
      // x->level[i].forward = update[i]->level[i].forward;
      // update[i]->level[i].forward = x;
      current = update[i];
      newNode->forward[i] = current->forward[i];
      current->forward[i] = newNode;
    }

    ++length;
  }

  ValueType &operator[](const KeyType &key) {
    ValueType value = find(key);
    if (value == tail->value)
      insert(key, 0);
    return find(key);
  }

  ValueType &find(const KeyType &key) {
    SkipListNode<KeyType, ValueType> *current = head;
    for (int i = currentLevel; i >= 0; --i) {
      while (current->forward[i]->key < key) {
        current = current->forward[i];
      }
    }
    current = current->forward[0];
    if (current->key == key)
      return current->value;
    return tail->value;
  }

  bool contains(const KeyType &key) { return find(key) != tail->value; }

private:
  std::mt19937 randomGenerator;
  std::uniform_int_distribution<> distribution;
};

int main() {
  SkipList<int, int> skipList;
  vector<int> keys(n);
  vector<int> values(n);

  for (int i = 0; i < n; ++i) {
    keys[i] = rand();
    values[i] = rand();
  }

  auto t1 = std::chrono::high_resolution_clock::now();
  for (int i = 0; i < n; ++i) {
    int key = keys[i], value = values[i];
    skipList[key] = value;
  }
  auto t2 = std::chrono::high_resolution_clock::now();

  auto duration =
      std::chrono::duration_cast<std::chrono::microseconds>(t2 - t1).count();
  std::cout << "Time cost: " << duration << " ms" << std::endl;

  return 0;
}